// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/check_duplicate_keys.h>
#include <pollux/functions/lib/lambda_function_util.h>
#include <pollux/functions/lib/rows_translation_util.h>
#include <pollux/vector/function_vector.h>

namespace kumo::pollux::functions {
    namespace {
        // See documentation at https://prestodb.io/docs/current/functions/map.html
        class TransformKeysFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_CHECK_EQ(args.size(), 2);

                // Flatten input map.
                exec::LocalDecodedVector mapDecoder(context, *args[0], rows);
                auto &decodedMap = *mapDecoder.get();

                auto flatMap = flattenMap(rows, args[0], decodedMap);

                std::vector<VectorPtr> lambdaArgs = {
                    flatMap->mapKeys(), flatMap->mapValues()
                };
                auto numKeys = flatMap->mapKeys()->size();

                SelectivityVector validRowsInReusedResult =
                        toElementRows<MapVector>(numKeys, rows, flatMap.get());

                VectorPtr transformedKeys;

                auto elementToTopLevelRows =
                        getElementToTopLevelRows(numKeys, rows, flatMap.get(), context.pool());
                const auto *rawElementToTopLevelRows =
                        elementToTopLevelRows->as<vector_size_t>();

                // Loop over lambda functions and apply these to keys of the map.
                // In most cases there will be only one function and the loop will run once.
                auto it = args[1]->as_unchecked<FunctionVector>()->iterator(&rows);
                while (auto entry = it.next()) {
                    auto keyRows =
                            toElementRows<MapVector>(numKeys, *entry.rows, flatMap.get());
                    auto wrapCapture = toWrapCapture<MapVector>(
                        numKeys, entry.callable, *entry.rows, flatMap);

                    entry.callable->apply(
                        keyRows,
                        &validRowsInReusedResult,
                        wrapCapture,
                        &context,
                        lambdaArgs,
                        elementToTopLevelRows,
                        &transformedKeys);

                    if (transformedKeys->may_have_nulls() ||
                        transformedKeys->may_have_nulls_recursive()) {
                        static const char *kNullKeyErrorMessage = "map key cannot be null";
                        static const char *kIndeterminateKeyErrorMessage =
                                "map key cannot be indeterminate";

                        keyRows.applyToSelected([&](vector_size_t keyRow) {
                            try {
                                POLLUX_USER_CHECK(
                                    !transformedKeys->is_null_at(keyRow), kNullKeyErrorMessage);

                                POLLUX_USER_CHECK(
                                    !transformedKeys->contains_null_at(keyRow),
                                    "{}: {}",
                                    kIndeterminateKeyErrorMessage,
                                    transformedKeys->toString(keyRow));
                            } catch (PolluxException &) {
                                context.setPolluxExceptionError(
                                    rawElementToTopLevelRows[keyRow], std::current_exception());
                            }
                        });
                    }
                }

                exec::LocalSelectivityVector remainingRows(context, rows);
                context.deselectErrors(*remainingRows);

                // Set nulls for rows not present in 'remainingRows'.
                BufferPtr newNulls = addNullsForUnselectedRows(flatMap, *remainingRows);

                auto localResult = std::make_shared<MapVector>(
                    flatMap->pool(),
                    outputType,
                    std::move(newNulls),
                    rows.end(),
                    flatMap->offsets(),
                    flatMap->sizes(),
                    transformedKeys,
                    flatMap->mapValues());

                checkDuplicateKeys(localResult, *remainingRows, context);

                context.moveOrCopyResult(localResult, rows, result);
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // map(K1, V), function(K1, V) -> K2 -> map(K2, V)
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("K1")
                    .typeVariable("K2")
                    .typeVariable("V")
                    .returnType("map(K2,V)")
                    .argumentType("map(K1,V)")
                    .argumentType("function(K1,V,K2)")
                    .build()
                };
            }
        };
    } // namespace

    /// transform_keys is null preserving for the map. But
    /// since an expr tree with a lambda depends on all named fields, including
    /// captures, a null in a capture does not automatically make a
    /// null result.

    POLLUX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
        udf_transform_keys,
        TransformKeysFunction::signatures(),
        exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
        std::make_unique<TransformKeysFunction>());
} // namespace kumo::pollux::functions
